1857G - Counting Graphs - CodeForces Solution

#include<bits/stdc++.h>
using namespace std;
#define CaseT int CaseT; cin >> CaseT; while(CaseT--)
#define endl '\n'
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> PII;
typedef pair<PII, int> PPI;
typedef pair<PII, PII> PPP;

const int INF = 0x3f3f3f3f;
const int N = 2e5 + 10;
const double eps = 1e-6;
const ll p = 998244353;

int n;
ll s;
struct EDGE {
    ll a, b, w;
}edges[N];
bool vis[N];
int f[N];
ll siz[N];

void init() {
    memset(vis, false, sizeof(vis));
    for (int i = 1; i <= n; i++) {
        f[i] = i;
        siz[i] = 1;
    }
}

bool cmp(EDGE &a, EDGE &b) {
    return a.w < b.w;
}

int find(int x) {
    if (f[x] != x) return f[x] = find(f[x]);
    return f[x];
}

ll qmi(ll a, ll k) {
    ll res = 1;
    while (k) {
        if (k & 1) res = res * a % p;
        k >>= 1;
        a = a * a % p;
    }
    return res;
}

void solve() {
    cin >> n >> s;
    init();
    for (int i = 1; i <= n - 1; i++) {
        ll a, b, w;
        cin >> edges[i].a >> edges[i].b >> edges[i].w;
    }

    sort(edges + 1, edges + n, cmp);

    ll res = 1;
    for (int i = 1; i <= n - 1; i++) {
        ll a = edges[i].a;
        ll b = edges[i].b;
        ll w = edges[i].w;
        a = find(a);
        b = find(b);

        if (a == b) continue;
        if (w == s) break;
        res = res * qmi(s - w + 1, siz[a] * siz[b] - 1) % p;
        siz[b] += siz[a];
        f[a] = b;
    }

    cout << res << endl;
}

int main() {
    cin.tie(0), cout.tie(0)->sync_with_stdio(false);
    CaseT
    solve();
    return 0;
}